import json
import yaml


def CreateModel(*args):
    yaml_file = args[0]
    with open(yaml_file) as f:
        model_def = yaml.safe_load(f)

    if model_def['type'] == args[1]:
        print(model_def)
        CreateCVModel()
    elif model_def['type'] == args[2]:
        print(model_def)
        CreateNLPModel()

    model_config = model_def.copy()
    del model_config['type']

    json_content = {
        'type': model_def['type'],
        'config': model_config
    }

    with open('config.json', 'w') as f:
        json.dump(json_content, f, indent=4)


def CreateCVModel():
    # 创建CV模型,返回模型对象
    print("CreateCVModel")


def CreateNLPModel():
    # 创建NLP模型,返回模型对象
    print("CreateNLPModel")


if __name__ == "__main__":
    CreateModel("CV_yaml.yaml", 'CV', 'NLP')
